/**********************************************************
read me:
    1. n can be up to 33,553,920(testing)
    2. haven't implemented scan algorithm
**********************************************************/

#include <stdio.h>
#include "cutil_math.h"

#define spaceDim 3

#define BLOCK_WIDTH 512

#define G 6.67384e-11f

#define MAX_NUM_STARS 33553920

__constant__ float softeningSquared =0; //the square of the softenting factor->EPS

//read n
unsigned int 
readN(void)
{

    unsigned int n = 0;
    
    if(scanf("%u", &n)!=1){
        fprintf(stderr, "Failed to read nunber of bodies-> n.\n");
            exit(1);
    }
    
    //check if n is not positive
    if(n<=0 || n > MAX_NUM_STARS){
        fprintf(stderr, "n is out of range.\n");
            exit(1);
    }
    
    getchar();
    
    return n;
}

//safely malloc
float4*
safe_malloc(int n){
    float4* ptr = NULL;
    size_t size = n*sizeof(float4);
    ptr = (float4*)malloc(size);
    if(ptr==NULL){
        fprintf(stderr, "malloc failed.\n");
            exit(1);
    }
    return ptr;
}
//read bodies
void
readBodies(unsigned int n, float4* body, float4* accel)
{
    int i =0;
    
    //read bodies
    for(i=0; i<n; i++){
        //read positions
        int j = 0;
        for(j=0; j<spaceDim; j++){
        
            switch(j){
                case 0:
                    if(scanf("%f", &body[i].x)!=1){
                        fprintf(stderr, "Failed to read positions.\n");
                            exit(1);
                    }
                    break;
                case 1:
                    if(scanf("%f", &body[i].y)!=1){
                        fprintf(stderr, "Failed to read positions.\n");
                            exit(1);
                    }
                    break;
                case 2:
                    if(scanf("%f", &body[i].z)!=1){
                        fprintf(stderr, "Failed to read positions.\n");
                            exit(1);
                    }
                    break;
            
            }
        
            
            
            getchar();
        }
        //read velocities
        j = 0;
        for(j=0; j<spaceDim; j++){
        
            switch(j){
                case 0:
                    if(scanf("%f", &accel[i].x)!=1){
                        fprintf(stderr, "Failed to read velocity.\n");
                            exit(1);
                    }
                    break;
                case 1:
                    if(scanf("%f", &accel[i].y)!=1){
                        fprintf(stderr, "Failed to read velocity.\n");
                            exit(1);
                    }
                    break;
                case 2:
                    if(scanf("%f", &accel[i].z)!=1){
                        fprintf(stderr, "Failed to read velocity.\n");
                            exit(1);
                    }
                    break;
            
            }
        }
        accel[i].w = 0.0f;//padding
        
        //read masses
        if(scanf("%f", &body[i].w)!=1){
                fprintf(stderr, "Failed to read masses.\n");
                    exit(1);
        }
            
        getchar();    
    }

}
//allocate memory on gpu
float4*
safe_cudaMalloc(const unsigned int n, float4* body, float4* velocity){
    float4* gpuPtr = NULL;
    size_t size = n*sizeof(float4);
    
    if(cudaMalloc(&gpuPtr, size)== cudaErrorMemoryAllocation){
        fprintf(stderr, "Failed to read masses.\n");
        free(body);
        free(velocity);
        exit(1);
    }
    return gpuPtr;

}


/****************************************************************************************************
for calculate accelerations
****************************************************************************************************/
//__device__ body-body
__device__ float3 
bodyBodyInteraction(float3 ai, float4 bi, float4 bj) 
{
    float3 r;

    // r_ij  [3 FLOPS]
    r.x = bj.x - bi.x;
    r.y = bj.y - bi.y;
    r.z = bj.z - bi.z;

    // distSqr = dot(r_ij, r_ij) + EPS^2  [6 FLOPS]
    float distSqr = r.x * r.x + r.y * r.y + r.z * r.z;
    distSqr += softeningSquared;

    // invDistCube =1/distSqr^(3/2)  [4 FLOPS (2 mul, 1 sqrt, 1 inv)]
    //float invDist = rsqrtf(distSqr);
    //float invDistCube =  invDist * invDist * invDist;

    float distSixth = distSqr * distSqr * distSqr;
    //for current case only-> soften factor is 0, when it is not 0 then don't need this if
    if(distSixth!=0){
        float invDistCube = 1.0f / sqrtf(distSixth);
        
        // s = m_j * invDistCube [1 FLOP]
        //float s = bj.w * invDistCube; //<- from Nvidia, it's wrong right? should be r_ij

        // a_i =  a_i + s * r_ij [6 FLOPS]
        ai.x += G*r.x * bj.w * invDistCube;
        ai.y += G*r.y * bj.w * invDistCube;
        ai.z += G*r.z * bj.w * invDistCube;
    
    }
    

    return ai;
}

//__device__ tile
__device__ float3
tile_calculation(float4 myPosition, float3 accel, float4* shPosition, unsigned int tileId, unsigned int GRID_WIDTH, bool rest, unsigned int restNum)// will this shPostition pointer/array stored on register or in global memory?
{
    int i =0;
    //last not copleted filled block/tile 
    if(rest&&tileId==(GRID_WIDTH-1)){
        for(i = 0; i < restNum; i++){
            accel = bodyBodyInteraction(accel, myPosition, shPosition[i]);
        }
    }
    //fully filled block/tile
    else{
        for(i = 0; i < blockDim.x; i++){
            accel = bodyBodyInteraction(accel, myPosition, shPosition[i]);
        }
    }
    
    return accel;
}

//__global__ do the calculation
__global__ void 
calculate_accel(void* devX, void* devA, int N, unsigned int GRID_WIDTH, bool rest, unsigned int restNum)
{
    __shared__ float4 shPosition[BLOCK_WIDTH];
    
    float4* globalX = (float4* )devX;
    float4* globalA = (float4*)devA;
    float4 myPosition;
    int i =0;
    int tile = 0;
    float3 acc;// = {0.0f, 0.0f, 0.0f};
    int gtid = blockIdx.x*blockDim.x + threadIdx.x;
    
    
    myPosition = globalX[gtid];

    acc = make_float3(globalA[gtid]);

    for(i = 0, tile = 0; i < N; i +=BLOCK_WIDTH, tile++){
        
        if(gtid < N){
            int idx = tile*blockDim.x + threadIdx.x;
            if(idx < N){
                shPosition[threadIdx.x] = globalX[idx];
            }
        }
        
        
        __syncthreads();
         
         if(gtid < N){
             acc = tile_calculation(myPosition, acc, shPosition, i, GRID_WIDTH, rest, restNum);
         }
        
        __syncthreads();
    }
    //save the result in global memory for the integration step
    if(gtid < N){
        float4 acc4 = {acc.x, acc.y, acc.z, 0.0f};
        globalA[gtid] = acc4;
    }
    
    
    
}
/****************************************************************************************************/

/****************************************************************************************************
update positions from derek
*****************************************************************************************************/
/*__global__ void updatePos(float4* pos, float3* vel, int n, float dt)
{
    int x = blockIdx.x*BLOCK_SIZE + threadIdx.x;

    if (x < n)
        pos[x] += make_float4(dt*vel[x]); // 3 flops
}*/
/*****************************************************************************************************/

int
main(void){
    //cpu
    float4* body_h = NULL;
    float4* velocity_h = NULL;
    
    unsigned int n = readN();
    bool rest = false;
    unsigned int restNum = 0;
    //gpu
    float4* body_d =NULL;
    float4* velocity_d = NULL;
    
    body_h = safe_malloc(n);
    velocity_h = safe_malloc(n);
    readBodies(n, body_h, velocity_h);
    
    //allocate memories on gpu
    body_d = safe_cudaMalloc(n, body_h, velocity_h);
    velocity_d = safe_cudaMalloc(n, body_h, velocity_h);
    
    //memcpy to device
       size_t size = n*sizeof(float4);
       cudaMemcpy(body_d, body_h, size, cudaMemcpyHostToDevice);
       cudaMemcpy(velocity_d, velocity_h, size, cudaMemcpyHostToDevice);
       
       //invoke kernel
       unsigned int GRID_WIDTH = n/BLOCK_WIDTH;
       restNum = n%BLOCK_WIDTH;
       if(restNum!=0){
           GRID_WIDTH += 1;
           rest = true;
       }
       calculate_accel<<<GRID_WIDTH, BLOCK_WIDTH>>>((void*)body_d, (void*)velocity_d, n, GRID_WIDTH, rest, restNum);
       
       //copy accels back to cpu
       cudaMemcpy(velocity_h, velocity_d, size, cudaMemcpyDeviceToHost);
       /**************************************************************************************************************/
       //test:
    int i =0;
    for(i=0;i<n;i++){
        //print body
        //printf("Body %d, mass %f:\n\n", i, body_h[i].w);
        
        //print positions
        //printf("positions:\n");
        printf("%9.2f", body_h[i].x);
        printf("%9.2f", body_h[i].y);
        printf("%9.2f", body_h[i].z);
        //printf("\n");
        
        //print velocity
        //printf("velocities:\n");
        
        printf("%9.2f", velocity_h[i].x);
        printf("%9.2f", velocity_h[i].y);
        printf("%9.2f", velocity_h[i].z);
        
        printf("\n\n");
    }
    //free gpu
    cudaFree(body_d);
    cudaFree(velocity_d);
    //free cpu
    free(body_h);
    free(velocity_h);
    return 0;
}